#include "kd_tree_heap.h"

#include <algorithm>
#include <iostream>
#include <stack>
#include <limits>

namespace Sirat { 

namespace Photon_Mapping {

Kd_Tree_Heap::Kd_Tree_Heap(unsigned int _max_photons) : max_photons(_max_photons) , size(0) {
	photons = new Photon*[_max_photons];
}

void print(Photon** _list,unsigned int _left,unsigned int _size) {
	
	for(unsigned int i = _left; i < _size +1; i++ ){ 
		std::cout << (_list[i])->get_position().x << " ";
	}
	std::cout << std::endl;
}

void Kd_Tree_Heap::construct_kd_tree(Photon** _list,unsigned int _left, unsigned int _right,unsigned int _heap_index,unsigned short _axis) {
// 	std::cout << "Constructing kd_tree / left: " <<  _left << " / right: " << _right << " Heap_index: " << _heap_index << std::endl;
	
	if(_left > _right) return;
	if(_axis != 3) _axis++;
	else _axis = 1;
	Photon::set_order(_axis);
	
	/* Getting a  pivot */
	
	unsigned int pivot = choose_pivot(_list,_left,_right,_axis);
	
// 	std::cout << "Pivot choose: " << pivot << std::endl;
	
	photons[_heap_index] = _list[pivot];
	size++;
		
	if(pivot != _left) { 
// 		std::cout << "<< Esquerda << " << std::endl;
		construct_kd_tree(_list,_left,pivot-1,2*_heap_index + 1,_axis);
	}
	
	if(pivot != _right) {
// 		std::cout << ">> Direita >> " << std::endl;
		construct_kd_tree(_list,pivot+1,_right,2*_heap_index + 2,_axis);
	}
}

unsigned int Kd_Tree_Heap::choose_pivot(Photon** _list,unsigned int _left,unsigned int _right,unsigned short _axis) {
	if(_left == _right) return _left;
	
	unsigned int size =  _right - _left;
	
// 	std::cout << "Choosing pivot / Size: " << size << std::endl;
	
	Photon** list = new Photon*[size+1];
	
	/* Copy list */
	unsigned int j = 0;
	for(unsigned int i = _left;i <= _right ; i++) {
		list[j] = _list[i];
		j++;
	}
	/* Sort and return the median */
// 	std::cout << "Printing list: " << std::endl;
// 	print(list,0,size);
	std::sort( list, list + size+1 , compare);
	
	for(unsigned int i = _left;i < _right+1; i++) {
		if((list[(size/2)]) == _list[i]) return i;
	}	
		
	//delete[] list;
	exit(1);
}

Photon* Kd_Tree_Heap::nearest_neighbour(const Math::Vector& _point) const {
	
	double root_distance = photons[0]->get_position().distance(_point);
	double minimal_distance = root_distance;
	Photon* root_photon = photons[0];
	Photon* nearest_photon = root_photon;
	
	// Tuple: (dimension,index,flag)
	std::stack< tuple<unsigned short,unsigned int,unsigned short> > s;
	/* Flags 
	0 = Go ahead
	1 = Look at left brother
	2 = Look at right brother
	*/
	s.push(make_tuple(1,0,0));
	while(!(s.empty())) {
		tuple<unsigned short,unsigned int,unsigned short> current = s.top();
		s.pop();
		
		/* Tuple values */
		const unsigned short dimension = current.get<0>();
		//Kd_Node* kd_node = current.get<1>();
		unsigned int index = current.get<1>();
		unsigned short flag = current.get<2>();
		const unsigned short next_dimension = (dimension != 3)?dimension + 1:1;
		
		const double value = _point[next_dimension];
		const double current_value = photons[index]->get_position()[next_dimension];
		
		switch(flag) {
			case 0: {
// 				Kd_Node* next_node = 0;
				int next_index = -1;
				int next_other_index = -1;
// 				Kd_Node* next_other_node = 0;
				unsigned short next_flag = 0;
				if(value <= current_value) {
					next_index = 2*index +1;
					next_other_index = 2*index +2;
					next_flag = 2;
				}
				else {
					next_index = 2*index +2;
					next_other_index = 2*index +1;
					next_flag = 1;
				}
				if(next_index != -1 && photons[next_index] != 0) {
					double next_distance = _point.distance(photons[next_index]->get_position());
					if(next_distance <= minimal_distance) {
						minimal_distance = next_distance;
						nearest_photon = photons[next_index];
					}
					s.push(make_tuple(dimension,index,next_flag));
					s.push(make_tuple(next_dimension,next_index,0));
				}
				else if(next_other_index != -1 && photons[next_other_index]) {
					double next_other_distance = _point.distance(photons[next_other_index]->get_position());
					if(next_other_distance <= minimal_distance) {
						minimal_distance = next_other_distance;
						nearest_photon = photons[next_other_index];
					}
					s.push(make_tuple(dimension,index,next_flag));
					s.push(make_tuple(next_dimension,next_other_index,0));
				}
				break;
			}
			/* Left brother! */
			case 1: {
				if(photons[2*index+1] != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double test_distance = _point.distance(photons[2*index+1]->get_position());
						if(test_distance < minimal_distance) {
							minimal_distance = test_distance;
							nearest_photon = photons[2*index+1];
						}
						s.push(make_tuple(next_dimension,2*index+1,0));
					}
				}
				break;
			}
			/* Right brother */
			case 2: {
				if(photons[2*index+2] != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double test_distance = _point.distance(photons[2*index+2]->get_position());
						if(test_distance < minimal_distance) {
							minimal_distance = test_distance;
							nearest_photon = photons[2*index+2];
						}
						s.push(make_tuple(next_dimension,2*index+2,0));
					}
				}
				break;
			}	
		}
	}
	return nearest_photon;
}

std::vector<tuple<Photon*,double> > Kd_Tree_Heap::k_nearest_neighbour(const Math::Vector& _point,const unsigned short _k) const {
	
	std::numeric_limits<double> double_limit;
	double minimal_distance = double_limit.infinity();
	
	double root_distance = photons[0]->get_position().distance(_point);
	Photon* root_photon = photons[0];

	
	std::vector<tuple<Photon*,double> > nearest;
	
	nearest.push_back(make_tuple(root_photon,root_distance));
	
	std::stack< tuple<unsigned short,unsigned int,unsigned short> > s;
	
	/* Flags 
	0 = Go ahead
	1 = Look at left brother
	2 = Look at right brother
	*/
	
	s.push(make_tuple(1,0,0));
	
	while(!(s.empty())) {
		tuple<unsigned short,unsigned int,unsigned short> current = s.top();
		s.pop();
		
		/* Tuple values */
		const unsigned short dimension = current.get<0>();
		//Kd_Node* kd_node = current.get<1>();
		unsigned int index = current.get<1>();
		unsigned short flag = current.get<2>();
		const unsigned short next_dimension = (dimension != 3)?dimension + 1:1;
		
		const double value = _point[next_dimension];
		const double current_value = photons[index]->get_position()[next_dimension];
	
		switch(flag) {
			case 0: {
				int next_index = -1;
				int next_other_index = -1;
// 				Kd_Node* next_node = 0;
// 				Kd_Node* next_other_node = 0;
				unsigned short next_flag = 0;
				
				if(value <= current_value) {
					next_index = 2*index + 1;
					next_other_index = 2*index + 2;
					next_flag = 2;
				}
				else {
					next_index = 2*index + 2;
					next_other_index = 2*index + 1;
					next_flag = 1;
				}
				if(next_index != -1 && photons[next_index] != 0) {
					double next_distance = _point.distance(photons[next_index]->get_position());
					Photon* next_photon = photons[next_index];
					if(next_distance <= minimal_distance) {
						insert_knn_list2(nearest,next_photon,next_distance,minimal_distance,_k);
					}
					s.push(make_tuple(dimension,index,next_flag));
					s.push(make_tuple(next_dimension,next_index,0));
				}
				else if(next_other_index != -1 && photons[next_other_index] != 0) {
					double next_distance = _point.distance(photons[next_other_index]->get_position());
					Photon* next_photon = photons[next_other_index];
					if(next_distance <= minimal_distance) insert_knn_list2(nearest,next_photon,next_distance,minimal_distance,_k);
					s.push(make_tuple(dimension,index,next_flag));
					s.push(make_tuple(next_dimension,next_other_index,0));
				}
				break;
			}
			/* Left brother! */
			case 1: {
				if(photons[2*index+1] != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double next_distance = _point.distance(photons[2*index+1]->get_position());
						Photon* next_photon = photons[2*index+1];
						if(next_distance <= minimal_distance) insert_knn_list2(nearest,next_photon,next_distance,minimal_distance,_k);
						s.push(make_tuple(next_dimension,2*index+1,0));
					}
				}
				break;
			}
			/* Right brother */
			case 2: {
				if(photons[2*index+2] != 0) {
					if(fabs(current_value - value) < minimal_distance) {
						double next_distance = _point.distance(photons[2*index+2]->get_position());
						Photon* next_photon = photons[2*index+2];
						if(next_distance <= minimal_distance) insert_knn_list2(nearest,next_photon,next_distance,minimal_distance,_k);
						s.push(make_tuple(next_dimension,2*index+2,0));
					}
				}
				break;
			}	
		}
	}
	return nearest;
}

/* This function tries to insert a possible nearest neighbour at this list */
void insert_knn_list2(std::vector<tuple<Photon*,double> >& nearest,Photon* next_photon,const double& next_distance,double& minimal_distance,const unsigned short _k) {
	if(nearest.size() == _k) {
		bool exist = false;
		/* Checking if this photon isnt at 'nearest' */
		for(unsigned int i = 0; i < nearest.size(); i++)
			if(nearest[i].get<0>() == next_photon) { exist = true; break; } 
		if(exist == false) 
			for(std::vector<tuple<Photon*,double> >::iterator i = nearest.begin() ; i != nearest.end() ; i++ ) {
			if( next_distance < (*i).get<1>() ) {
				nearest.insert(i,make_tuple(next_photon,next_distance));
				nearest.erase(nearest.end()-1);
				minimal_distance = (*(nearest.end()-1)).get<1>();
				break;
			}
			}
	} else {
		bool insert = false;
		bool exist = false;
		/* Checking if this photon isnt at 'nearest' */
		for(unsigned int i = 0; i < nearest.size(); i++)
			if(nearest[i].get<0>() == next_photon) { exist = true; break; } 
		if(exist == false) {
			for(std::vector<tuple<Photon*,double> >::iterator i = nearest.begin() ; i != nearest.end() ; i++ ) {
				if( next_distance < (*i).get<1>() ) {
					nearest.insert(i,make_tuple(next_photon,next_distance));
					insert = true;
					break;
				}
			}
			if(insert == false) nearest.insert(nearest.end(),make_tuple(next_photon,next_distance));
		}
	}
}

}

}
